import cv2
import numpy as np
import os
import datetime
class PoseVisualizer:
    """Pose可视化类，用于在原图上同时显示标准位置和当前位置"""
    
    def __init__(self):
        # 颜色定义
        self.colors = {
            'left_top': (0, 255, 0),    # 绿色
            'left_0': (255, 0, 255),    # 紫色
            'left_bottom': (255, 0, 0),  # 蓝色
            'bottom_0': (0, 128, 128),   # 深青色
            'right_bottom': (0, 255, 255), # 黄色
            'right_0': (255, 255, 0),    # 青色
            'right_top': (0, 0, 255),    # 红色
            'top_0': (128, 0, 128),      # 深紫色
        }
        # 定义连接顺序（多边形连接顺序）
        self.connect_order = [
            'left_top', 'left_0', 'left_bottom',
            'bottom_0', 'right_bottom', 'right_0',
            'right_top', 'top_0', 'left_top'  # 闭合多边形
        ]
        
    def visualize_pose(self, pose_dict, st_pose_dict, image_path, timestamp,output_dir="results_pose_visualization"):
        """
        在原图上可视化标准位置和当前位置
        
        Args:
            pose_dict: 当前位置字典，如 {'left_top': (980.0, 180.0), ...}
            st_pose_dict: 标准位置字典
            image_path: 原始图像路径（str）或图片数据（numpy.ndarray）
            output_dir: 输出目录
            
        Returns:
            保存可视化结果的图像路径
        """
        os.makedirs(output_dir, exist_ok=True)
        
        # 处理输入图像
        if isinstance(image_path, np.ndarray):
            # 直接使用传入的图片数据
            image = image_path.copy()
            # 生成时间戳作为文件名
            filename = f"pose_{timestamp}.png"
        else:
            # 从路径读取图片
            image = cv2.imread(image_path)
            if image is None:
                print(f"错误: 无法读取图像 {image_path}")
                return None
            # 使用原始文件名
            filename = os.path.basename(image_path)
            filename = f"pose_{filename}"
        
        output_path = os.path.join(output_dir, filename)
        
        # 1. 绘制标准位置（绿色实线）
        image = self._draw_reference_pose(image, st_pose_dict)
        
        # 2. 绘制当前位置（红色实线）
        image = self._draw_current_pose(image, pose_dict)
        
        cv2.imwrite(output_path, image)
        return output_path
    
    def _draw_current_pose(self, image, pose_dict):
        """绘制当前位置（红色实线）"""
        if pose_dict is None:
            return image
        
        points = []
        # 收集所有有效点
        for name in self.connect_order:
            if name in pose_dict and pose_dict[name]:
                x, y = pose_dict[name]
                points.append((int(x), int(y)))
        
        # 绘制多边形（红色实线）
        if len(points) >= 2:
            cv2.polylines(image, [np.array(points)], isClosed=True, color=(0, 0, 255), thickness=2)
        
        # 绘制点和标签
        for name in pose_dict:
            if pose_dict[name]:
                x, y = pose_dict[name]
                point = (int(x), int(y))
                if name in self.colors:
                    color = self.colors[name]
                    cv2.circle(image, point, 5, color, -1)
                    cv2.putText(image, name, (point[0]+10, point[1]), 
                               cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1)
        
        return image
    
    def _draw_reference_pose(self, image, st_pose_dict):
        """绘制标准位置（绿色实线）"""
        if st_pose_dict is None:
            return image
        
        points = []
        # 收集所有有效点
        for name in self.connect_order:
            if name in st_pose_dict and st_pose_dict[name]:
                x, y = st_pose_dict[name]
                points.append((int(x), int(y)))
        
        # 绘制多边形（绿色实线）
        if len(points) >= 2:
            cv2.polylines(image, [np.array(points)], isClosed=True, color=(0, 255, 0), thickness=2)
        
        # 绘制参考点（小圆圈）
        for name in st_pose_dict:
            if st_pose_dict[name]:
                x, y = st_pose_dict[name]
                point = (int(x), int(y))
                cv2.circle(image, point, 3, (192, 192, 192), -1)
        
        return image

# 使用示例
if __name__ == "__main__":
    # 示例数据
    st_pose_dict = {
        'left_top': (100, 100),
        'left_0': (100, 200),
        'left_bottom': (100, 300),
        'bottom_0': (200, 300),
        'right_bottom': (300, 300),
        'right_0': (300, 200),
        'right_top': (300, 100),
        'top_0': (200, 100)
    }
    
    current_pose_dict = {
        'left_top': (120, 110),
        'left_0': (110, 210),
        'left_bottom': (105, 305),
        'bottom_0': (220, 310),
        'right_bottom': (320, 310),
        'right_0': (310, 210),
        'right_top': (290, 110)
        # 缺少'top_0'点测试跳过功能
    }
    
    visualizer = PoseVisualizer()
    result_path = visualizer.visualize_pose(
        current_pose_dict,
        st_pose_dict,
        "./data/Color1/440.png",
        "output_results"
    )
    print(f"可视化结果保存至: {result_path}")